# coding=utf-8

import math
import multiprocessing
import platform
import queue
import sys
import os
import time
import traceback
import logging
import whisper
import zhconv
from functools import partial
from datetime import datetime
from pydub import AudioSegment
from threading import Lock
from faster_whisper import WhisperModel

logging.basicConfig(
    level=logging.DEBUG,
    format='%(asctime)s %(filename)s[line:%(lineno)d] %(levelname)s %(message)s',
    datefmt='%Y-%m-%d %H:%M:%S',
    filename='output.log',
    filemode='a'
)
logger = logging.getLogger(__name__)

CURRENT_FILE_PATH = os.path.abspath(__file__)
CURRENT_DIR_PATH = os.path.dirname(CURRENT_FILE_PATH)
MODEL_PATH = os.path.join(os.path.join(CURRENT_DIR_PATH, "sources"), "models")  # 模型所在路径
VIDEO_EXTENSIONS = ['.mp4', '.avi', '.mov', '.mkv']  # 添加你需要的视频文件扩展名


class ExecuteZiMu:
    """
    提取字幕
    """

    def __init__(self, video_path, output_path):
        super().__init__()
        self.video_path = video_path
        self.output_path = output_path
        self.model = whisper.load_model(name="base", download_root=MODEL_PATH)
        # self.model = whisper.load_model(name="medium", download_root=MODEL_PATH)

        # 音频相关
        self.segment_length = 1 * 60 * 1000  # 每个音频的长度（单位：毫秒）

    @staticmethod
    def video_to_audio(video_path):
        """
        从视频中提取音频
        :return:
        """
        audio_path = f"{os.path.splitext(video_path)[0]}.wav"  # "path/to/audio.wav"

        # 使用Pydub加载视频文件
        video = AudioSegment.from_file(video_path)

        # 提取视频中的音频部分
        audio = video.set_channels(1)  # 将音频转换为单声道（可选）
        audio.export(audio_path, format="wav")  # 导出音频文件

        return audio_path

    def extract(self, video_path):
        audio_path_list = []  # 所有音频路径
        # 总音频
        content = ""

        try:
            st = time.perf_counter()

            # 从视频中提取音频
            audio_path = self.video_to_audio(video_path)
            logger.info(f"audio_path: {audio_path}\tvideo_to_audio cost time: {time.perf_counter() - st}")

            # 音频切割
            audio = AudioSegment.from_file(audio_path)
            audio_path_list.append(audio_path)

            # 音频总时长
            total_duration = len(audio)

            # 计算需要切割的音频段数
            # num_segments = total_duration // segment_length
            num_segments = math.ceil(total_duration / self.segment_length)  # 向上取整，音频能够获取完整
            logger.info(
                f"total_duration: {total_duration}\tsegment_length: {self.segment_length}\tnum_segments:{num_segments}")

            # 切割后的音频保存位置
            split_audio_path = os.path.splitext(audio_path)[0]

            # 循环切割音频
            for i in range(num_segments):
                # 计算当前音频的起始时间和结束时间
                start_time = i * self.segment_length
                end_time = (i + 1) * self.segment_length
                if end_time > total_duration:
                    end_time = total_duration

                logger.info(f"start_time: {start_time}\tend_time: {end_time}")
                segment = audio[start_time:end_time]

                # 保存切割后的音频为新文件
                output_file = f"{split_audio_path}_{i}.wav"
                segment.export(output_file, format="wav")
                audio_path_list.append(output_file)

                # 语音转文本识别
                result = self.model.transcribe(output_file, language='zh', verbose=False)

                # 繁体转简体
                segments = result['segments']

                for seg in segments:
                    # 繁体转简体
                    res = zhconv.convert(seg['text'], "zh-cn")
                    content += res + ","

                content += "\n\n\n"

            logger.info(f"content: {content}")
            logger.info(f"cost time: {time.perf_counter() - st}")
        except Exception as e:
            logger.error(f"ExecuteZiMu ERROR: {traceback.format_exc()}")
        finally:
            # 移除所生成的音频
            for path in audio_path_list:
                os.remove(path)
                logger.info(f"{path} 移除成功！")

        return content

    def extract_for_file(self):
        """
        单个视频
        :return:
        """
        content = self.extract(self.video_path)

        txt_path = os.path.join(self.output_path, f"{os.path.splitext(self.video_path)[0]}.txt")

        # 写入视频名称到txt文件
        ExecuteZiMu.write_to_file(txt_path, content=content)

    def extract_for_dirs(self):
        """
        从目录中提取，并保留原有目录结构
        :return:
        """
        # 确保目标目录存在
        if not os.path.exists(self.output_path):
            os.makedirs(self.output_path)

        for dirpath, dirnames, filenames in os.walk(self.video_path):
            # 获取相对于源目录的路径
            rel_dir = os.path.relpath(dirpath, self.video_path)
            # 构造目标目录的路径
            target_dir = os.path.join(self.output_path, rel_dir)

            # 如果目标子目录不存在，则创建它
            if not os.path.exists(target_dir):
                os.makedirs(target_dir)

            for filename in filenames:
                # 检查文件扩展名是否是视频文件
                if any(filename.lower().endswith(ext) for ext in VIDEO_EXTENSIONS):
                    # 构造视频文件的完整路径
                    video_path = os.path.join(dirpath, filename)
                    try:
                        # 构造txt文件的完整路径（在目标目录中）
                        txt_path = os.path.join(target_dir, f"{os.path.splitext(filename)[0]}.txt")

                        # 提取字幕
                        content = self.extract(video_path)

                        # 写入视频名称到txt文件
                        ExecuteZiMu.write_to_file(txt_path, content=content)

                        logger.info(f"Created {txt_path} for {video_path} success!")
                    except:
                        logger.error(f"extract_error: {traceback.format_exc()}\tvideo_path: {video_path}")

    @staticmethod
    def write_to_file(path, content):
        with open(path, 'w', encoding='utf-8') as txt_file:
            txt_file.write(content + '\n')

    def dispatch(self):
        """
        分发
        :return:
        """
        if not os.path.exists(self.video_path):
            logger.error(f"{self.video_path} not exits.")
            return

        if os.path.isfile(self.video_path):
            self.extract_for_file()
        else:
            self.extract_for_dirs()


def main():
    """
    base: 00:03:09 cost time: 37.55908260000069
    medium: cost time: 133.84059049999996
    :return:
    """
    video_path = r"xxx.mp4"
    output_path = r""

    video_path = r"G:\xxx\sources\videos"
    output_path = r"G:\xxx\sources\videos\notebook"

    ezm = ExecuteZiMu(video_path=video_path, output_path=output_path)
    ezm.dispatch()



if __name__ == '__main__':
    main()
